package com.dahe.ysxt.vector.vectorDB;

import com.tencent.tcvectordb.client.VectorDBClient;
import com.tencent.tcvectordb.model.Collection;
import com.tencent.tcvectordb.model.Database;
import com.tencent.tcvectordb.model.DocField;
import com.tencent.tcvectordb.model.Document;
import com.tencent.tcvectordb.model.param.database.ConnectParam;
import com.tencent.tcvectordb.model.param.dml.InsertParam;
import com.tencent.tcvectordb.model.param.enums.ReadConsistencyEnum;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class VectorDBUtils {
    public static VectorDBClient buildVectorDBClient() {
        System.out.println("\tvdb_url: " + System.getProperty("vdb_url"));
        System.out.println("\tvdb_key: " + System.getProperty("vdb_key"));
        System.setProperty("vdb_url", "http://lb-rmfsvvoh-84vj9k7fcdan3o2x.clb.ap-beijing.tencentclb.com:40000");
        System.setProperty("vdb_key", "vepGVsPr9ddIZgmvhaw5T0C5NOrKdSaee2MH4UA1");
        ConnectParam connectParam = ConnectParam.newBuilder()
                .withUrl(System.getProperty("vdb_url"))
                .withUsername("root")
                .withKey(System.getProperty("vdb_key"))
                .withTimeout(30)
                .build();
        VectorDBClient client = new VectorDBClient(connectParam, ReadConsistencyEnum.EVENTUAL_CONSISTENCY);
        return client;
    }

    /**
     * 向指定数据库集合中插入文档
     * @param client
     * @param dbName 数据库名
     * @param collectionName 集合名
     * @param documentList 文档集合
     * @throws InterruptedException
     */
    public static void upsertData(VectorDBClient client, String dbName, String collectionName, List<Document> documentList) throws InterruptedException {
        Database database = client.database(dbName);
        Collection collection = database.describeCollection(collectionName);
        InsertParam insertParam = InsertParam.newBuilder().addAllDocument(documentList).withBuildIndex(true).build();
        collection.upsert(insertParam);
        // notice：upsert 操作可用会有延迟
        Thread.sleep(1000 * 5);
    }
}
